package pipeline

/**
 * 带睡眠功能的 workPool
 */

import (
	"context"
	"fmt"
	"sync"
)

type TaskFunc func(args ...interface{})

// 可以把参数传到函数内
type Task struct {
	F    TaskFunc
	Args []interface{}
}

type WorkPool struct {
	pool        chan *Task
	workerCount int

	// stop相关
	stopCtx        context.Context
	stopCancelFunc context.CancelFunc
	wg             sync.WaitGroup
}

func (t *Task) Execute() {
	t.F(t.Args...)
}

func NewWorkPool(workerCount, poolLen int) *WorkPool {
	return &WorkPool{
		workerCount: workerCount,
		pool:        make(chan *Task, poolLen),
	}
}

func (w *WorkPool) PushTask(t *Task) {
	w.pool <- t
}

func (w *WorkPool) PushTaskFunc(f TaskFunc, args ...interface{}) {
	w.pool <- &Task{
		F:    f,
		Args: args,
	}
}

func (w *WorkPool) work(i int) {
	for {
		select {
		case <-w.stopCtx.Done():
			w.wg.Done()
			return
		case t := <-w.pool:
			t.Execute()
		}
	}
}

func (w *WorkPool) Start() *WorkPool {
	fmt.Printf("workpool run %d worker\n", w.workerCount)
	w.wg.Add(w.workerCount)
	w.stopCtx, w.stopCancelFunc = context.WithCancel(context.Background())
	for i := 0; i < w.workerCount; i++ {
		go w.work(i)
	}
	return w
}

func (w *WorkPool) Stop() {
	fmt.Printf("workpool stop %d worker\n", w.workerCount)
	w.stopCancelFunc()
	w.wg.Wait()
}
